import networkx as nx
import matplotlib.pyplot as plt
from collections import deque
import numpy as np
from networkx.classes import Graph, graph
from networkx.drawing import circular_layout


def generate_directed_ba_graph(nodes=20, edges=2):
    """生成有向Barabási-Albert无标度网络"""
    # 先生成无向图，然后随机分配方向
    undirected_G = nx.barabasi_albert_graph(nodes, edges)
    G = nx.DiGraph()
    G.add_nodes_from(undirected_G.nodes())
    for u, v in undirected_G.edges():
        if np.random.rand() < 0.5:
            G.add_edge(u, v)
        else:
            G.add_edge(v, u)
    return G


# 3. 将有向图转换为二分图（节点拆分）
def convert_to_bipartite(digraph):
    """将有向图转换为二分图：每个节点拆分为入节点和出节点"""
    B = nx.Graph()

    # 创建左部节点（出节点）和右部节点（入节点）
    left_nodes = [f"{node}_out" for node in digraph.nodes()]
    right_nodes = [f"{node}_in" for node in digraph.nodes()]

    B.add_nodes_from(left_nodes, bipartite=0)
    B.add_nodes_from(right_nodes, bipartite=1)

    # 添加边：原图中的有向边(u, v)对应二分图中u_out到v_in的边
    for u, v in digraph.edges():
        B.add_edge(f"{u}_out", f"{v}_in")

    return B, left_nodes, right_nodes


# 4. 实现Hopcroft-Karp算法
def hopcroft_karp(bipartite_graph, left_nodes):
    """Hopcroft-Karp算法实现二分图最大匹配"""
    # 初始化距离和匹配字典
    distance = {}
    match = {node: None for node in bipartite_graph.nodes()}

    # BFS：查找是否存在增广路径
    def bfs():
        queue = deque()
        for node in left_nodes:
            if match[node] is None:
                distance[node] = 0
                queue.append(node)
            else:
                distance[node] = float('inf')

        distance[None] = float('inf')

        while queue:
            node = queue.popleft()
            if distance[node] < distance[None]:
                for neighbor in bipartite_graph.neighbors(node):
                    next_node = match.get(neighbor, None)
                    if distance[next_node] == float('inf'):
                        distance[next_node] = distance[node] + 1
                        queue.append(next_node)
        return distance[None] != float('inf')

    # DFS：查找增广路径
    def dfs(node):
        if node is None:
            return True
        for neighbor in bipartite_graph.neighbors(node):
            next_node = match.get(neighbor, None)
            if distance[next_node] == distance[node] + 1:
                if dfs(next_node):
                    match[neighbor] = node
                    match[node] = neighbor
                    return True
        distance[node] = float('inf')
        return False

    # 主算法循环
    matching_size = 0
    while bfs():
        for node in left_nodes:
            if match[node] is None and dfs(node):
                matching_size += 1
    return match, matching_size


# 5. 可视化原始有向图和匹配结果
def visualize_directed_graph(digraph, matched_edges, title):
    """可视化有向图并高亮显示匹配边"""
    plt.figure(figsize=(10, 10))
    # pos = nx.random_layout(digraph, seed=42)
    pos = circular_layout(digraph)

    # 绘制所有节点和边
    nx.draw_networkx_nodes(digraph, pos, node_size=200, node_color='salmon')
    nx.draw_networkx_edges(digraph, pos,
                           edge_color='gray',
                           width=1.5,
                           alpha=0.6,
                           arrowstyle='->',
                           arrowsize=15)
    nx.draw_networkx_labels(digraph, pos, font_size=10)

    # 高亮显示匹配边
    nx.draw_networkx_edges(digraph, pos,
                           edgelist=matched_edges,
                           edge_color='red',
                           width=2.0,
                           arrowstyle='->',
                           arrowsize=18)

    # 标记匹配节点
    matched_nodes = set()
    for u, v in matched_edges:
        # matched_nodes.add(u)
        matched_nodes.add(v)
    # 标记未匹配节点
    unmatched_nodes = set(digraph.nodes()) - matched_nodes
    nx.draw_networkx_nodes(digraph, pos,
                           nodelist=unmatched_nodes,
                           node_size=400,
                           node_color='lightblue')

    plt.title(title, fontsize=15)
    plt.axis('off')
    plt.tight_layout()
    plt.show()


def visualize_bipartite_with_matching(G, left_nodes, right_nodes):
    """可视化二分图及匹配结果"""
    plt.figure(figsize=(8, 8))

    # 设置节点位置
    pos = {}
    # 左部节点位置（左侧一列）
    for i, node in enumerate(left_nodes):
        pos[node] = (1, -i * 2)
    # 右部节点位置（右侧一列）
    for i, node in enumerate(right_nodes):
        pos[node] = (4, -i * 2)

    # 绘制所有节点
    nx.draw_networkx_nodes(G, pos, nodelist=left_nodes,
                           node_color='skyblue',
                           node_size=800,
                           edgecolors='black')
    nx.draw_networkx_nodes(G, pos, nodelist=right_nodes,
                           node_color='lightgreen',
                           node_size=800,
                           edgecolors='black')

    # 添加节点标签
    labels = {}
    for node in G.nodes():
        labels[node] = node

    nx.draw_networkx_labels(G, pos, labels, font_size=12, font_weight='bold')

    # 绘制所有边（灰色）
    nx.draw_networkx_edges(G, pos,
                           edgelist=G.edges(),
                           edge_color='lightgray',
                           width=2,
                           alpha=0.7)
    # 添加图例
    plt.text(1, 6, "左部节点", fontsize=12, color='darkblue')
    plt.text(4, 6, "右部节点", fontsize=12, color='darkgreen')

    # 设置坐标轴
    plt.xlim(-1, 5)
    plt.ylim(-len(left_nodes) * 2 - 3, 2)
    plt.axis('off')
    plt.tight_layout()
    plt.show()


def example1_graph():
    exa1_graph = Graph(directed=True)
    left_nodes = ['A', 'B', 'C', 'D', 'E', 'F', 'G']
    right_nodes = ['1', '2', '3', '4', '5', '6', '7']
    exa1_graph.add_nodes_from(left_nodes)
    exa1_graph.add_nodes_from(right_nodes)
    edges = [('A', '1'), ('A', '4'),
             ('B', '3'), ('B', '6'), ('B', '7'),
             ('C', '2'), ('C', '4'), ('C', '5'),
             ('D', '2'), ('D', '7'),
             ('E', '5'), ('E', '6'), ('E', '7'),
             ('F', '3'), ('F', '6'),
             ('G', '6'), ('G', '7')]
    exa1_graph.add_edges_from(edges)
    return exa1_graph, left_nodes, right_nodes


# 主程序
if __name__ == "__main__":
    # plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
    # plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
    #
    # exa1_graph, exa1_left_nodes, exa1_right_nodes = example1_graph()
    # visualize_bipartite_with_matching(exa1_graph, exa1_left_nodes, exa1_right_nodes)
    # exa1_match, exa1_matching_size = hopcroft_karp(exa1_graph, exa1_left_nodes)


    # 生成有向ER网络
    er_nodes_size, er_probability = 6, 0.1
    er_digraph = nx.erdos_renyi_graph(er_nodes_size, er_probability, directed=True, seed=42)
    print(f"ER有向图: {len(er_digraph.nodes())}节点, {len(er_digraph.edges())}边")

    # 转换为二分图
    er_bipartite, er_left, er_right = convert_to_bipartite(er_digraph)

    # 计算最大匹配
    er_match, er_size = hopcroft_karp(er_bipartite, er_left)
    print(f"ER网络最大匹配边数: {er_size}")
    print(f"驱动节点数量：{er_nodes_size - er_size}")

    # 提取原图中的匹配边
    er_matched_edges = []
    for out_node, in_node in er_match.items():
        if out_node and in_node and out_node.endswith("_out") and in_node.endswith("_in"):
            u = int(out_node.split("_")[0])
            v = int(in_node.split("_")[0])
            er_matched_edges.append((u, v))

    # 可视化结果
    visualize_directed_graph(er_digraph, er_matched_edges, "Directed ER Network with Maximum Matching")

